/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2020-2022. All rights reserved.
 * Description: HCOM 工具类
 * Author: lilianlin
 * Create: 2020-12-05
 */

#ifndef HCOM__OP_UTILS_H
#define HCOM__OP_UTILS_H

#include <map>
#include <string>
#include "external/hccl/hccl_types.h"
#include "graph/op_desc.h"
#include "graph/utils/node_utils.h"
#include "platform/platform_info.h"
#include "base.h"


#include "ops_kernel_builder_base.h"
#include "op_hcom_comm.h"
#include "hccl_common.h"
#include "hcom_common.h"
#include "comm_base_pub.h"

#include "graph/compute_graph.h"

namespace hccl {
const std::string AUTOTUNE_HCCL_OPS_LIB_NAME = "ops_kernel_info_hccl_gradtune";
const std::string HCCL_OPS_LIB_NAME = "ops_kernel_info_hccl";

const std::string HCOM_ATTR_RANK_SIZE = "rank_size";
const std::string HCOM_ATTR_SHAPE = "shape";

const std::map<ge::DataType, HcclDataType> HCOM_DATA_TYPE_MAP = {
    {ge::DT_FLOAT, HCCL_DATA_TYPE_FP32},
    {ge::DT_FLOAT16, HCCL_DATA_TYPE_FP16},
    {ge::DT_INT8, HCCL_DATA_TYPE_INT8},
    {ge::DT_INT32, HCCL_DATA_TYPE_INT32},
    {ge::DT_INT16, HCCL_DATA_TYPE_INT16},
    {ge::DT_INT64, HCCL_DATA_TYPE_INT64},
    {ge::DT_UINT64, HCCL_DATA_TYPE_UINT64},
    {ge::DT_UINT8, HCCL_DATA_TYPE_UINT8},
    {ge::DT_UINT16, HCCL_DATA_TYPE_UINT16},
    {ge::DT_UINT32, HCCL_DATA_TYPE_UINT32},
    {ge::DT_DOUBLE, HCCL_DATA_TYPE_FP64},
    {ge::DT_BF16, HCCL_DATA_TYPE_BFP16}
};

const std::map<std::string, HcclReduceOp> HCOM_REDUCE_TYPE_MAP{
    {"sum", HCCL_REDUCE_SUM},
    {"prod", HCCL_REDUCE_PROD},
    {"max", HCCL_REDUCE_MAX},
    {"min", HCCL_REDUCE_MIN},
};

enum class CreateDir {
        HCCL_DIR_NUM_NEG = -1,
        HCCL_DIR_NUM_ZEO = 0,
        HCCL_DIR_NUM_ONE = 1,
        HCCL_DIR_NUM_TWO = 2,
        HCCL_DIR_NUM_TRE = 3
    };

struct OffSendRecvInfo {
    // 存放数据长度和偏移长度
    std::vector<u64> sendLength;
    // 存放数据个数和偏移个数
    std::vector<u64> sendCounts;
};
constexpr int DEFAULT_SERVER_NUM = 2;

// Ge适配的类
constexpr u32 SEND_RECEIVE_TASK_NUM = 20;
constexpr u32 OP_DEFAULT_TASK_NUM = 245;
constexpr u32 ONE_DEVICE_TASK_NUM = 5;
constexpr u32 AICORE_MIN_CLEAR_ZEOR_SIZE = 32;
constexpr u32 DFX_DEFAULT_TASK_NUM = 16; // 2个计数 + 2个其它dfx + 2个memcpy_async + 预留10个
constexpr u32 DFX_PADDING_TASK_NUM = 4;
constexpr u32 MASTER_STREAM_EVENT_NUM = 3;
constexpr u32 SLAVE_STREAM_EVENT_NUM = 2;
constexpr u32 COM_STEP_NUM = 2;
constexpr s32 SERVER_NUM_ONE = 1;
constexpr s32 SERVER_NUM_EIGHT = 8;
constexpr u32 ALLREDUCE_DEFAULT_COM_STEP = 19;  // allgather + reducescatter
constexpr u32 ALLGATHER_DEFAULT_COM_STEP = 9;  // 5个通信 + 4个barrier
constexpr u32 REDUCESCATTER_DEFAULT_COM_STEP = 11;  // 6个通信 + 5个barrier
constexpr u32 ALLTOALL_DEFAULT_COM_STEP = 14; // alltoall taskNum, 每个对端通信的最大task数量(图模式, NA+pairwise)
constexpr u32 TASK_NUM_DEVICE_FOUR = 4;
constexpr u32 TASK_NUM_DEVICE_TWO = 2;
constexpr s32 TASK_NUM_DEVICE_ONE = 1;
constexpr u32 ALG_8P_RING_COMM_STEP = 7;
constexpr u32 PIPLINE_STREAM_EVENT_NUM = 2;

constexpr u32 TASK_MAX_NUM_DEV_TYPE_V80 = 2042;
constexpr u32 TASK_MAX_NUM_DEV_TYPE_V71 = 4061;

class HcomOpUtils {
public:
    HcomOpUtils() = default;
    ~HcomOpUtils() = default;
    static HcclResult GetReduction(const ge::OpDescPtr &opDescPtr, HcclReduceOp &reduction);
    static HcclResult GetRoot(const ge::OpDescPtr &opDescPtr, int32_t &root);
    static HcclResult GetRankSize(const ge::OpDescPtr &opDescPtr, int32_t &rankSize);
    static HcclResult GetGroup(const ge::OpDescPtr &opDescPtr, std::string &group);
    static HcclResult GetDataType(const ge::OpDescPtr &opDescPtr, HcclDataType &dataType);
    static HcclResult GetDataType(const ge::OpDescPtr &opDescPtr, std::string &dataType);
    static HcclResult GetSrcRank(const ge::OpDescPtr &opDescPtr, int32_t &srcRank);
    static HcclResult GetDestRank(const ge::OpDescPtr &opDescPtr, int32_t &destRank);
    static HcclResult GetSrTag(const ge::OpDescPtr &opDescPtr, int32_t &srTag);
    static HcclResult TransformDataType(const ge::DataType geDataType, HcclDataType &hcclDataType);
    static HcclResult GetAllInputsTensorMemSize(const ge::OpDescPtr &opDescPtr, uint64_t &tensorSize);
    static HcclResult GetAllOutputsTensorMemSize(const ge::OpDescPtr &opDescPtr, uint64_t &tensorSize);
    static HcclResult GetAllInputsTensorOriginSize(const ge::OpDescPtr &opDescPtr, uint64_t &tensorSize);
    static HcclResult GetTensorMemSize(const ge::GeTensorDesc &tensorDesc, uint64_t &memSize);
    static HcclResult GetTensorOriginSize(const ge::GeTensorDesc &tensorDesc, uint64_t &size);
    static HcclResult GetPathFromEnv(const char* environment, std::string &fusionPath);
    static bool HcomOpIsSupportedBool(const std::string &opType);
    static HcclResult ConversionOpDataType(const ge::OpDescPtr &op, const std::string &opType, HcclDataType &dataType);
    static HcclResult GetAlltoAllvcStagedScratchMemSize(ge::Node& node, const int64_t &hcomComm, const string &sGroup,
        u32 rankSize, u64 &getMemSize);
    static HcclResult GetAlltoAllCountMatrix(const ge::OpDescPtr &op, std::vector<int64_t> &sendCountMatrix);
    static HcclResult GetAlltoAllCountMatrix(ge::Node& node, std::vector<int64_t> &sendCountMatrix);
    static HcclResult GetAlltoAllDataType(const ge::OpDescPtr &op, HcclDataType& sendType, HcclDataType& recvType);
    static HcclResult GetConstInputAcrossGraph(const ge::GeTensor* &tensor, u32 index, ge::Node& node);
    static HcclResult GetVectorFromTensor(const ge::GeTensor* tensor, std::vector<int64_t>& vector);
    static HcclResult GetRankId(const int64_t &hcomComm, const string &sGroup, u32 &rankId);
    static HcclResult GetDevId(const int64_t &hcomComm, const string &sGroup, s32 &devId);
    static HcclResult CheckAlltoAllvcRank(const ge::Node &node, const int64_t &hcomComm, const string &sGroup);
    static bool IsNeedCalTaskNum(const std::string &sCollectiveType);
    static HcclResult GetDfxTaskNum(const std::string &sCollectiveType, u32 &taskNum);
    static HcclResult GetToSlaveStreamTaskNum(const std::string &sCollectiveType, u64 streamNum, u64 piplineSliceNum,
        u32 &taskNum);
    static HcclResult GetToMasterStreamTaskNum(const std::string &sCollectiveType, u32 &taskNum);
    static HcclResult GetDefaultAlgoLevel0Module(s32 deviceNumPerServer, AlgTypeLevel0 &algType);
    static HcclResult GetDefaultAlgoLevel1(s32 serverNum, AlgTypeLevel1 &algType);
    static HcclResult GetAlgoLevel1(s32 serverNum, std::string &opType, AlgTypeLevel1 &algType);
    static HcclResult SplitHcclOpType(const std::string &algoConfig, const std::string &opType,
        std::string &specificAlgoConfig);
    static HcclResult GetAlgType(const ge::Node &node, s32 deviceNumPerServer, s32 serverNum,
        std::string opType, AlgType &algType);
    static HcclResult GetDevAndSerNumFromRankTable(s32 &deviceNumPerServer, s32 &serverNum,
        bool &multiModuleDiffDeviceNumMode);
    static HcclResult GetGroupFromOpDesc(const ge::OpDescPtr &op, std::string &sGroup);
    static HcclResult GetDeviceAndServerNum(const ge::Node &node, s32 &deviceNumPerServer, s32 &serverNum,
        bool &multiModuleDiffDeviceNumMode);
    static HcclResult GetBetweenServersStep(s32 serverNum, u32 &commStep);
    static HcclResult GetIntraComTaskNum(const std::string &sCollectiveType, s32 deviceNumPerServer,
        u64 streamNum, const AlgType &algType, u32 &taskNum, u64 totalSize);
    static HcclResult GetInterComTaskNum(const std::string &sCollectiveType, s32 serverNum, s32 deviceNumPerServer,
        DevType devType, u32 &taskNum);
    static HcclResult GetCombineComTaskNum(const std::string &sCollectiveType, s32 serverNum, s32 deviceNumPerServer,
        u32 &intraTaskNum, u32 &interTaskNum);
    static HcclResult GetTensorNum(const ge::Node &node, const std::string &sCollectiveType, u32 &tensorNum);
    static HcclResult GetTensorSize(const ge::GeTensorDesc &tensorDesc, int64_t &size);
    static HcclResult GetAllTensorSize(const ge::OpDescPtr &op, u32 tensorNum, std::vector<int64_t> &tensorSize);
    static HcclResult GetTaskNumFromCrackSize(const ge::Node &node, u32 tensorNum, u32 &taskNum);
    static HcclResult GetTensorCleanTaskNum(const ge::Node &node, const std::string &sCollectiveType, u32 &taskNum);
    static HcclResult GetAllReduceScratchMemSize(const string &sGroup, const ge::Node& node,
        s32 serverNum, s32 rankSize, const u64 count, const HcclDataType dataType, u64 &getMemSize);
    static HcclResult GetModuleInfo(const std::vector<RankInfo_t> &rankList, bool &multiModuleDiffDeviceNumMode);
#ifndef HCOM_EXECUTOR
    static HcclResult CreateFusionConfigVersion(std::string &configVersion);
    static HcclResult GetFileNameFromPath(std::string &Path, std::string &fusionFile);
#endif
};
}

#endif // end HCOM__OP_UTILS_H
